import cv2 as cv
import sys
import os
from python_ai.common.xcommon import *
import matplotlib.pyplot as plt
import numpy as np
import time
import datetime


# def my_show_img(img, title="no title", trans=None, **kwargs):
#     global spn
#     spn += 1
#     plt.subplot(spr, spc, spn)
#     if trans is not None:
#         img = trans(img)
#     plt.imshow(img, **kwargs)
#     plt.axis('off')
#     plt.title(title)


# img_dir = '../../../../../large_data/CV2/lesson/Day06'
# img_name = 'hammer.jpg'
img_dir = '../../../../../large_data/pic'
# img_name = 'DSC05039.JPG'
# img_name = 'DSC05022_1.JPG'
# img_name = 'dog.jpg'
# img_name = 'dog_bird.jpg'
# img_name = 'football.jpg'
img_name = 'messi5.jpg'
img_path = os.path.join(img_dir, img_name)

# spr = 2
# spc = 3
# spn = 0
# plt.figure(figsize=[9, 6])

sep('load')
img = cv.imread(img_path, cv.IMREAD_GRAYSCALE)
print('original shape', img.shape)
H, W = img.shape
H2 = (H - 1) // 2
W2 = (W - 1) // 2
# my_show_img(img, 'original', cmap='gray')


def on_change(event):
    pass


title = 'fft bars'
# cv.namedWindow(title, cv.WINDOW_NORMAL)
cv.namedWindow(title)
cv.imshow(title, np.zeros((50, 1000), dtype=np.uint8))
cv.createTrackbar('range', title, 30, W2, on_change)
cv.createTrackbar('Is HPF', title, 1, 1, on_change)

img = cv.imread(img_path, cv.IMREAD_GRAYSCALE)
img_cpy = img.copy()

rate = 750 / W
img_show = cv.resize(img_cpy, None, fx=rate, fy=rate, interpolation=cv.INTER_CUBIC)
cv.imshow('img_show', img_show)

FPS = 25
INTERVAL = 1000 // FPS
QUIT_SET = set([ord('q'), ord('Q'), 27])
while True:

    f = np.fft.fft2(img)
    eps = 1e-8
    # f_show = np.log(abs(f) + eps)
    # my_show_img(f_show, 'f_show', cmap='gray')

    f_shift = np.fft.fftshift(f)
    f_shift_show = np.log(abs(f_shift) + eps)
    f_shift_show = cv.normalize(f_shift_show, None, 0., 255., cv.NORM_MINMAX, dtype=cv.CV_8UC1)
    # my_show_img(f_shift_show, 'f_shift_show', cmap='gray')

    HALF = cv.getTrackbarPos('range', title)
    is_hpf = cv.getTrackbarPos('Is HPF', title)
    # Conclusion: center in shifted dft are low frequency components
    if is_hpf:  # high pass
        f_shift[H2 - HALF:H2 + HALF, W2 - HALF:W2 + HALF] = 0  # erase low, i.e. erase center
    else:  # low pass
        mask = np.zeros((H, W), dtype=np.float32)
        mask[H2 - HALF:H2 + HALF, W2 - HALF:W2 + HALF] = 1  # reserve low, i.e. reserve center
        f_shift *= mask
    # f_shift_show = np.log(abs(f_shift) + eps)
    # my_show_img(f_shift_show, 'f_shift_show clipped', cmap='gray')

    if_ = np.fft.ifftshift(f_shift)
    # if_show = np.log(abs(if_) + eps)
    # my_show_img(if_show, 'if_show', cmap='gray')

    img_ = np.fft.ifft2(if_)
    img_ = abs(img_)
    # img_ = np.log(img_ + 1e-20)
    # my_show_img(img_, 'filtered', cmap='gray')

    f_shift_show = cv.cvtColor(f_shift_show, cv.COLOR_GRAY2BGR)
    cv.rectangle(f_shift_show, (W2 - HALF, H2 - HALF), (W2 + HALF, H2 + HALF), (0, 255, 0), 2)
    if is_hpf:
        cv.line(f_shift_show, (W2 - HALF, H2 - HALF), (W2 + HALF, H2 + HALF), (0, 255, 0), 2)
    else:
        cv.line(f_shift_show, (0, 0), (W2 - HALF, H2 - HALF), (0, 255, 0), 2)
        cv.line(f_shift_show, (W2 + HALF, H2 + HALF), (W, H), (0, 255, 0), 2)

    min = img_.min()
    max = img_.max()
    img_ -= min
    img_ /= max - min
    img_ *= 255.
    img_ = img_.astype(np.uint8)
    img_ = cv.cvtColor(img_, cv.COLOR_GRAY2BGR)

    img_all = np.concatenate((f_shift_show, img_), axis=1)
    cv.line(img_all, (W, 0), (W, H), (0, 255, 0), 2)

    img_all = cv.resize(img_all, None, fx=rate, fy=rate, interpolation=cv.INTER_CUBIC)
    cv.imshow('FFT', img_all)

    k = cv.waitKey(INTERVAL) & 0xFF
    if k in QUIT_SET:
        break

cv.destroyAllWindows()
